from typing import List, Optional, Dict, Any, Tuple
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStore
from src.ingestion.vector_store import VectorStoreManager
import numpy as np
from abc import ABC, abstractmethod

class BaseRetriever(ABC):
    """检索器基类"""
    @abstractmethod
    def retrieve(self, query: str, **kwargs) -> List[Document]:
        """检索接口"""
        pass

class VectorDBRetriever(BaseRetriever):
    """基于向量数据库的检索器"""
    def __init__(
        self,
        vector_store_manager: VectorStoreManager,
        top_k: int = 4,
        score_threshold: Optional[float] = None
    ):
        """
        初始化向量数据库检索器
        :param vector_store_manager: 向量存储管理器实例
        :param top_k: 默认返回结果数量
        :param score_threshold: 分数阈值，低于此值的结果将被过滤
        """
        self.vector_store_manager = vector_store_manager
        self.top_k = top_k
        self.score_threshold = score_threshold

    def retrieve(self, query: str,** kwargs) -> List[Document]:
        """
        执行检索
        :param query: 查询文本
        :param kwargs: 额外参数 (top_k, score_threshold等)
        :return: 检索到的文档列表
        """
        # 处理参数
        top_k = kwargs.get('top_k', self.top_k)
        score_threshold = kwargs.get('score_threshold', self.score_threshold)

        # 执行带分数的相似性检索
        results_with_scores = self.vector_store_manager.similarity_search_with_score(
            query=query,
            k=top_k
        )

        # 过滤分数阈值
        filtered_results = []
        for doc, score in results_with_scores:
            if score_threshold is None or score <= score_threshold:
                # 将分数添加到文档元数据
                doc.metadata['score'] = score
                filtered_results.append(doc)

        return filtered_results

class HybridRetriever(BaseRetriever):
    """混合检索器（向量检索+关键词检索）"""
    def __init__(
        self,
        vector_store_manager: VectorStoreManager,
        top_k_vector: int = 3,
        top_k_keyword: int = 2,
        weight_vector: float = 0.7,
        weight_keyword: float = 0.3
    ):
        """
        初始化混合检索器
        :param vector_store_manager: 向量存储管理器实例
        :param top_k_vector: 向量检索返回结果数量
        :param top_k_keyword: 关键词检索返回结果数量
        :param weight_vector: 向量检索权重
        :param weight_keyword: 关键词检索权重
        """
        self.vector_store_manager = vector_store_manager
        self.top_k_vector = top_k_vector
        self.top_k_keyword = top_k_keyword
        self.weight_vector = weight_vector
        self.weight_keyword = weight_keyword

    def _keyword_search(self, query: str, k: int = 2) -> List[Tuple[Document, float]]:
        """简单关键词检索（基于元数据和文本中的关键词匹配）"""
        # 此处实现简化版关键词检索
        # 实际应用中可使用BM25等算法
        vector_results = self.vector_store_manager.similarity_search_with_score(query, k=k*3)
        keyword_scores = []

        # 提取查询关键词
        query_keywords = set(query.lower().split())

        for doc, _ in vector_results:
            # 在文档内容和元数据中查找关键词
            content = doc.page_content.lower()
            metadata_str = ' '.join([str(v).lower() for v in doc.metadata.values()])
            combined_text = f"{content} {metadata_str}"

            # 计算关键词匹配得分
            matched_keywords = query_keywords.intersection(set(combined_text.split()))
            score = len(matched_keywords) / len(query_keywords) if query_keywords else 0
            keyword_scores.append((doc, score))

        # 按得分排序并取前k个
        keyword_scores.sort(key=lambda x: x[1], reverse=True)
        return keyword_scores[:k]

    def retrieve(self, query: str, **kwargs) -> List[Document]:
        """执行混合检索"""
        # 获取参数
        top_k_vector = kwargs.get('top_k_vector', self.top_k_vector)
        top_k_keyword = kwargs.get('top_k_keyword', self.top_k_keyword)

        # 向量检索
        vector_results = self.vector_store_manager.similarity_search_with_score(
            query=query,
            k=top_k_vector
        )
        # 关键词检索
        keyword_results = self._keyword_search(query, k=top_k_keyword)

        # 合并结果并去重
        doc_id_map = {} 
        for doc, score in vector_results:
            # 使用文档元数据中的id作为唯一标识，确保同一文档不会被重复计数
            doc_id = doc.metadata.get('id', str(id(doc)))  # 转换为字符串以确保一致性
            doc.metadata['vector_score'] = score
            doc.metadata['retrieval_type'] = 'vector'
            doc_id_map[doc_id] = (doc, score * self.weight_vector)

        for doc, score in keyword_results:
            # 使用文档元数据中的id作为唯一标识
            doc_id = doc.metadata.get('id', str(id(doc)))  # 转换为字符串以确保一致性
            doc.metadata['keyword_score'] = score
            doc.metadata['retrieval_type'] = 'keyword'
            if doc_id in doc_id_map:
                # 已存在，合并分数
                existing_doc, existing_score = doc_id_map[doc_id]
                combined_score = existing_score + (score * self.weight_keyword)
                # 合并元数据
                for key, value in doc.metadata.items():
                    if key not in existing_doc.metadata:
                        existing_doc.metadata[key] = value
                doc_id_map[doc_id] = (existing_doc, combined_score)
            else:
                doc_id_map[doc_id] = (doc, score * self.weight_keyword)

        # 按合并分数排序并返回
        combined_results = sorted(doc_id_map.values(), key=lambda x: x[1], reverse=True)
        final_docs = [doc for doc, _ in combined_results]

        # 添加最终合并分数
        for i, (doc, score) in enumerate(combined_results):
            final_docs[i].metadata['combined_score'] = score

        return final_docs

class RetrievalManager:
    """检索管理器，统一管理不同检索策略"""
    def __init__(
        self,
        vector_store_manager: VectorStoreManager,
        retrieval_strategy: str = "vector",
        **kwargs
    ):
        """
        初始化检索管理器
        :param vector_store_manager: 向量存储管理器实例
        :param retrieval_strategy: 检索策略 (vector/hybrid)
        :param kwargs: 检索器参数
        """
        self.vector_store_manager = vector_store_manager
        self.retrieval_strategy = retrieval_strategy

        # 根据策略初始化检索器
        if retrieval_strategy == "vector":
            self.retriever = VectorDBRetriever(
                vector_store_manager=vector_store_manager,
                **kwargs
            )
        elif retrieval_strategy == "hybrid":
            self.retriever = HybridRetriever(
                vector_store_manager=vector_store_manager,
                **kwargs
            )
        else:
            raise ValueError(f"不支持的检索策略: {retrieval_strategy}")

    def retrieve(self, query: str,** kwargs) -> List[Document]:
        """执行检索"""
        return self.retriever.retrieve(query, **kwargs)

    def switch_strategy(self, strategy: str,** kwargs) -> None:
        """切换检索策略"""
        self.retrieval_strategy = strategy
        # 重新初始化检索器
        if strategy == "vector":
            self.retriever = VectorDBRetriever(
                vector_store_manager=self.vector_store_manager,
                **kwargs
            )
        elif strategy == "hybrid":
            self.retriever = HybridRetriever(
                vector_store_manager=self.vector_store_manager,
                **kwargs
            )
        else:
            raise ValueError(f"不支持的检索策略: {strategy}")

if __name__ == "__main__":
    # 示例用法
    from src.ingestion.vector_store import VectorStoreManager

    try:
        # 初始化向量存储
        vector_store = VectorStoreManager(
            embedding_model="huggingface",
            vector_store_type="chroma",
            persist_directory="./vector_db_chroma"
        )

        # 初始化检索管理器（向量检索策略）
        retrieval_manager = RetrievalManager(
            vector_store_manager=vector_store,
            retrieval_strategy="vector",
            top_k=3
        )

        # 测试向量检索
        query = "什么是RAG技术?"
        print(f"\n向量检索查询: {query}")
        results = retrieval_manager.retrieve(query)
        for i, doc in enumerate(results):
            print(f"结果 {i+1} (分数: {doc.metadata.get('score'):.4f}):")
            print(f"来源: {doc.metadata.get('source')}")
            print(f"内容: {doc.page_content[:150]}...\n")

        # 切换到混合检索策略
        retrieval_manager.switch_strategy("hybrid", top_k_vector=3, top_k_keyword=2)

        # 测试混合检索
        print(f"混合检索查询: {query}")
        hybrid_results = retrieval_manager.retrieve(query)
        for i, doc in enumerate(hybrid_results):
            print(f"结果 {i+1} (综合分数: {doc.metadata.get('combined_score'):.4f}):")
            print(f"检索类型: {doc.metadata.get('retrieval_type')}")
            print(f"向量分数: {doc.metadata.get('vector_score', 'N/A'):.4f}")
            print(f"关键词分数: {doc.metadata.get('keyword_score', 'N/A'):.4f}")
            print(f"内容: {doc.page_content[:150]}...\n")

    except Exception as e:
        print(f"检索时出错: {str(e)}")